#include "PhotonMap.h"
#include <trax.hpp>
#include "atMaterial.h"
#include "PointLight.h"
#include "atMathConstant.h"
#include "atMaterial.h"
#include "PhotonBucket.h"
#include "RayTracer.h"
#include "Shader.h"
#include "atAtomicReg.h"
#include "atUtil.h"
#include "MemoryManager.h"

#define USE_PHT_BKT 1
#define STORE_PHT_ATOMIC 1

#define EMIT_PHT_NUM 1000
#define MAX_PHT_BOUNCE 1
#define MAX_TOT_PHT (EMIT_PHT_NUM*(MAX_PHT_BOUNCE+1)) // max total number of photons
//#define NUM_THREADS 4 // hard code style, should be careful

//#define PHT_STORE_PER_THREAD (MAX_TOT_PHT / NUM_THREADS)
#define PHT_CHECKED 1000 // photon checked in z-curve rad estimate
//#define PHT_CHECKED MAX_TOT_PHT
#define PHT_COUNT 500 	// maximum photon take in count

#if 0
    #define MIN_BOUNCE 1
    #define MAX_BOUNCE (MIN_BOUNCE+1) // for testing
#else
    #define MIN_BOUNCE 0
    #define MAX_BOUNCE MAX_PHT_BOUNCE
#endif

#define PRINT_SORT_PROCESS 1
#define TEST_PHTM_SORT 0

namespace AtomTrace
{
	int GetMaxPhotonCapacity()
	{ return (MAX_TOT_PHT * sizeof(Photon) / 4); }
    
    int GetMaxNumberOfPhotons()
    { return MAX_TOT_PHT; }
    
    int GetPhotonMapEndAddr()
    {
        return (GetPhotonMapAddr() + GetMaxPhotonCapacity());
    }
    
    int GetSingleEmittedPhotons()
    { return EMIT_PHT_NUM; }
    
    // -------------- Photon Member functions ------------------
	void Photon::SaveToMemory(int addr)
	{
		flux.SaveToGM(addr);
		pRay._o.SaveToGM(addr + sizeof(flux)/4);
		pRay._dir.SaveToGM(addr + (sizeof(flux)+sizeof(pRay._o)) / 4 );
	}
    
	void Photon::LoadFromMemory(int addr)
	{
		flux.LoadFromGM(addr);
		pRay._o.LoadFromGM(addr+(sizeof(flux) / 4));
		pRay._dir.LoadFromGM(addr + (sizeof(flux)+sizeof(pRay._o)) / 4);
	}

	void Photon::LoadFromMemoryGivenIndex(int index)
	{
		int addr = index*(sizeof(Photon) / 4) + GetPhotonMapAddr()+ 1;
		this->LoadFromMemory(addr);
	}

    void Photon::DebugPrint() const
    {
        printf("flux: ");
        for (int i=0; i<3; i++) {
            printf("%f, ", flux[i]);
        }
        printf("\npos: ");
        for (int i=0; i<3; i++) {
            printf("%f, ", pRay._o[i]);
        }
        printf("\ndir: ");
        for (int i=0; i<3; i++) {
            printf("%f, ", pRay._dir[i]);
        }
        printf("\n");
    }

    // ---------------------------------------------------------------------
    
	Photon LoadPhotonFromGM(int addr)
	{
		Photon pht;
		//int addr = index*(sizeof(Photon) / 4) + GetPhotonMapAddr()+ 1;
		pht.flux.LoadFromGM(addr);
		pht.pRay._o.LoadFromGM(addr+(sizeof(pht.flux) / 4));
		pht.pRay._dir.LoadFromGM(addr + (sizeof(pht.flux)+sizeof(pht.pRay._o)) / 4);
		return pht;
	}
    
    Photon LoadPhotonFromGMGivenIndex(int index)
    {
        int addr = index*(sizeof(Photon) / 4) + GetPhotonMapAddr()+ 1;
        return LoadPhotonFromGM(addr);
    }

	void PhotonSwap(Photon & pht1, Photon & pht2, int & zid1, int &zid2, const int addr1, const int addr2)
	{
		pht1.SaveToMemory(addr2);
		pht2.SaveToMemory(addr1);

		Photon pht3 = pht1;
		pht1 = pht2;
		pht2 = pht3;

		int t=zid1;
		zid1 = zid2;
		zid2 = t;
	}

    // -------------------- Photon Map Member functions ------------------------
    // ===============================================================================

    PhotonMap::PhotonMap():
    		pm_s_addr(-1),
    		num_pht(-1),
    		pht_size(sizeof(Photon) / 4), // size in word
    		_tid(-1),
            s_mat_addr(-1),
    		sorted(false)
    {
    	pm_s_addr = GetPhotonMapAddr();
    	int baddr = GetBVH(); // bvh root node contain everything
    	sbox.pmin.LoadFromGM(baddr);
    	sbox.pmax.LoadFromGM(baddr+3);
    }

	void PhotonMap::Compute(const Scene & scene, const int threadId)
	{
		//if(_tid < 0)
		_tid = threadId;
		/*
		 * while(numpht < emi_photon)
		 * 	Generate new photon
		 * 	whlie(bounce++<maxbounce)
		 * 		TracePhoton();
		 * 		if(hit){
		 * 			if(hit.diffuse){
		 * 				store(photon);
		 * 			}
		 * 			new dir;
		 * 		}else
		 * 			break;
		 */

		int maxBounce = MAX_PHT_BOUNCE;
		PointLight pl;
		pl.LoadFromGM(scene.startLightAddr + 0 * 3); // only consider the first light for debugging

		Color oneLInten = pl.GetIntensity();
		float totIntens = oneLInten.Gray();
        if(_tid == 0)
            printf("Light Intens: %f\n", totIntens);
		//const float fac = totIntens * PI_F * 4.0f / static_cast<float>(EMIT_PHT_NUM);
		//Color ptInten = ( oneLInten / oneLInten.Gray() ) * fac;
        Color ptInten = pl.GetIntensity() * 4.0f * PI_F;
        //ptInten = pl.GetIntensity() * PI_F * 40.0f / static_cast<float>(EMIT_PHT_NUM);

		Photon pht;
		PhotonBucket pbck;

		const int pbk_s_addr = pm_s_addr;
		int usedSpot = 0;
		int cur_pbk_addr = pbk_s_addr + 1; // the first word store the number of photons in the current bucket in memory
        int cur_pht_offset;
		int start_pbk_addr = cur_pbk_addr;
		int pbsize = pbck.GetLocalStackSize();
        int phtSize = sizeof(Photon) / 4;
        //printf("PhotonBucket size in word: %d\n", pbsize);
        int numThreads = trax_nthreads();
        int num_pht_t = EMIT_PHT_NUM / numThreads;
        
        for(int phtid = 0; phtid < num_pht_t; phtid++)
		{
			//float p=100.f*(phtid+1)/EMIT_PHT_NUM; // percentage display
			//printf("\rPhotonPass %5.2f%%",p);
            //printf("\nPhton %d", phtid);
			pht.Set(pl.RandomPhoton());
			pht.flux = ptInten;
			int bounce = 0;
            int bkCount = 0;

			while(bounce<maxBounce && pht.flux.Gray() > 0.00001f)
			{
                bounce++;
				HitRecord hit;
				hit.Init();
				bool isHit = false;
				isHit = RayTracer::TraceBVH(pht.pRay, hit);
				if(isHit)
				{
                    hit.UseTriangleNormal();
					int matAddr = scene.startMatAddr + hit.matID * 25;
					hit.matIllum = loadi(matAddr, 0);
					if(Material::IsPhotonSurface(hit.matIllum)){
                        
                        pht.pRay._o = hit.P;
                        //if(phtid == 0)
                        //{ pht.DebugPrint(); }
#if USE_PHT_BKT
						if( bounce > MIN_BOUNCE && bounce <= MAX_BOUNCE )
                        //if(bounce != 2 && bounce > MIN_BOUNCE)
                        {
                            if( !pbck.StorePhoton(pht) ) // if not success store into local stack photon bucket, it means it is full
                            {
    #if STORE_PHT_ATOMIC
                                pbck.StoreIntoMemory();
                                pbck.Clear();
                                pbck.StorePhoton(pht);
                            
    #else
                                printf("Not using atomic add\n");
                                pbck.StoreIntoMemory(cur_pbk_addr); // store into memory and calculate next memory region for current thread
                                bkCount++;
                                cur_pbk_addr += pbsize;
                                usedSpot += pbck.GetCurrent();
                                pbck.Clear();
                                pbck.StorePhoton(pht);
    #endif
                            }
                        }
#else
						cur_pht_offset = atomicinc(REG_PHT_ST_ID);
						pht.SaveToMemory(start_pbk_addr + cur_pht_offset * phtSize);
#endif
					}
                    float p = trax_rand();
                    if(p < 0.1f) break; // naive absorption
					Material::BouncePhoton(pht.pRay, pht.flux, hit, scene, ptInten);
				}else
				{
					break;
				}
			}

		}
        
        // finally store the rest photon into memory
#if USE_PHT_BKT
    #if STORE_PHT_ATOMIC
        pbck.StoreNonEmptyIntoMem();
        pbck.Clear();

    #else
        pbck.StoreNonEmptyIntoMem(cur_pbk_addr);
        usedSpot += pbck.GetCurrent();
        cur_pbk_addr += pbck.GetCurrentStackSize();
        pbck.Clear();
    #endif
#endif
        
        barrier(REG_BARRIER);
        if(threadId == 0){
#if USE_PHT_BKT
    #if STORE_PHT_ATOMIC
            cur_pht_offset = atomicadd(REG_PHT_ST_ID, 0); // make sure get the correct number of photons
        	storei(cur_pht_offset, pbk_s_addr, 0);
        	printf("Using atomicinc photon_bucket\nTotal photon number: %d\n", cur_pht_offset);
    #else
            storei(usedSpot, pbk_s_addr, 0);
            printf("Using photon_bucket without atomic\nTotal photon number: %d\n", usedSpot);
    #endif
#else
            storei(cur_pht_offset + 1, pbk_s_addr, 0);
        	printf("No photon bucket\nTotal photon number: %d\n", cur_pht_offset+1);
#endif
            Photon testPht = LoadPhotonFromGM(start_pbk_addr);
            printf("First photon:\n");
            testPht.DebugPrint();
        }
        barrier(REG_BARRIER);

        // num_pht is in thread stack so it must be load from memory after tid==0 done storage
        num_pht = loadi(pbk_s_addr, 0);
#if TEST_PHTM_SORT
        if(!sorted)
        {
            if(_tid == 0)
                printf("photon map sorting...\n");
            this->Sort();
            
            if(_tid==0)
                printf("Sorting done\n");
            
            if(_tid == 0)
            {
				Photon testPht = LoadPhotonFromGM(start_pbk_addr);
				printf("First photon after sorting:\n");
				testPht.DebugPrint();
            }
        }
        barrier(REG_BARRIER);
#endif
	}

	Photon PhotonMap::LoadPhotonGivenID(int index)
	{
		int addr = this->PhotonAddr(index);
		return LoadPhotonFromGM(addr);
	}

	int PhotonMap::GetZcurveIndexGivenPhotonId(int index, Photon & pht)
	{
		int addr = this->PhotonAddr(index);
		pht = LoadPhotonFromGM(addr);
		int res;
		GetZcurveIndex(pht.pRay._o, res);
		return res;
	}

	// assume photon map is ready, so num_pht and pm_s_addr must be valid.
	void PhotonMap::Sort()
	{
		if(sorted) return;
		pm_s_addr = GetPhotonMapAddr();

		// Single thread sort
		if(_tid == 0)
		{
#if 1
			// TRING: Quick sort
			const int ssize = 256;
			int stack[ssize];
			int top = -1;
			int l = 0;
			int r = num_pht-1;

			stack[++top] = l;
			stack[++top] = r;
			int iter = 0;
			while(top >= 0)
			{
				r = stack[top--];
				l = stack[top--];
				int p;
				// partition
				{
					Photon phtR;
					int zidR = GetZcurveIndexGivenPhotonId(r, phtR);
					int addrR = this->PhotonAddr(r);

					int i=l-1;
					Photon phtj, phti;
					int zidI;// = GetZcurveIndexGivenPhotonId(i, phti);
					int addrI;// = this->PhotonAddr(i);

					for(int j=l; j<r; j++)
					{
						int zidJ = GetZcurveIndexGivenPhotonId(j, phtj);

						int addrJ = this->PhotonAddr(j);

						if(zidJ < zidR)
						{
							i++;
							zidI = GetZcurveIndexGivenPhotonId(i, phti);
							addrI = this->PhotonAddr(i);
							PhotonSwap(phti, phtj, zidI, zidJ, addrI, addrJ);
						}
					}
					i++;
					zidI = GetZcurveIndexGivenPhotonId(i, phti);
					addrI = this->PhotonAddr(i);
					PhotonSwap(phti, phtR, zidI, zidR, addrI, addrR);

					p = i;
				}

				//printf("Pivot: %i\n", p);

				if(p-1 > l){
					stack[++top] = l;
					stack[++top] = p-1;
				}

				if(p+1 < r){
					stack[++top] = p+1;
					stack[++top] = r;
				}
				//printf("top: %i\n", top);
				if(top >= ssize)
				{
					printf("********** Error, stack full in quick sort ************\n");
				}
				iter++;
            }
				//if(iter > 10) break;
                //  QUICK SORT
#else
                // DOING: First try: bubble sort (or insert sort?)
                Photon pht1;
                Photon pht2;
                for (int i=0; i<num_pht; i++)
                {
                    if(i%500 == 0)
                        printf("main pht %d\n", i);
                    int addr1 = this->PhotonAddr(i);
                    pht1.LoadFromMemory(addr1);
                    int zid1=0;
                    GetZcurveIndex(pht1.pRay._o, zid1);
                    for(int j=i+1; j<num_pht; j++)
                    {
                        int addr2 = this->PhotonAddr(j);
                        pht2.LoadFromMemory(addr2);
                        int zid2 = 0;
                        GetZcurveIndex(pht2.pRay._o, zid2);
                        
                        if(zid2 < zid1) // swap
                        {
                            pht1.SaveToMemory(addr2);
                            pht2.SaveToMemory(addr1);
                            pht1 = pht2;
                            zid1 = zid2;
                            
                        }
                    }
                }

#endif
		}
		barrier(REG_BARRIER);
		sorted = true;
	}

	bool PhotonMap::BruteForceRadEst(HitRecord & hit, const float r, Color & res, float &omax_r2, float &nnei)
	{
        if(s_mat_addr < 0)
            s_mat_addr = GetMaterials();
        int matAddr = s_mat_addr + hit.matID * 25;
        
        Color diffuse;
        hit.matIllum = loadi(matAddr, 0);
        if(hit.matIllum == 3 || hit.matIllum == 6)
        {
            //printf("Non photon surface\n");
        	res.Set(0.0f, 0.0f, 0.0f); // ignore non-photon surface so far
        	return true;
        }
        diffuse.LoadFromGM(matAddr+4);
        
		// TODO: Brute force estimate
		PhotonBucket pbk;
		int pbk_cap = pbk.GetTotalPhotonCapacity();
		int delti = 0;
		int pht_size = sizeof(Photon) / 4; // size in word
		Color sumColor = 0.0f;
		Vector3 sumDir = 0.0f; // incoming direction
        nnei = 0;
		int neiCount = 0;
        float r2 = r * r;
        omax_r2 = 0.0f;
        float &maxR2 = omax_r2;
		for(int i=0; i<num_pht && neiCount < PHT_COUNT; i+=delti)
		{
			// using stack to increase locality
			int addr;
			if(i+pbk_cap < num_pht)
			{
				delti = pbk_cap;
				addr = pm_s_addr+1 +i*pht_size;
				pbk.LoadPhotonsFromGM( addr );
			}else
			{
				delti = num_pht - i;
				addr = pm_s_addr+1 + i*pht_size;
				pbk.LoadPhotonsFromGM( addr, delti);
			}
			// -------------------------------

			int pbk_cur = pbk.GetCurrent();
			
			for(int j=0; j<pbk_cur; j++)
			{
				const Vector3 & pos = pbk[j].pRay._o;
                float tr2 = (pos-hit.P).LengthSquare();
				if( tr2 < r2 && (pbk[j].pRay._dir.Dot(hit.GetTriangleNormal()) < 0.0f))
				{
                    if(tr2 > maxR2) maxR2 = tr2;
					sumColor += pbk[j].flux;// / (tr2+1.0f);
					sumDir += pbk[j].pRay._dir;
					neiCount++;
				}
			}
		}
        if(maxR2 < 0.001f) maxR2 = r2;
        //printf("neighbor count: %d\n", neiCount);

		// average sum
        //float inv_area = 1.0f / (PI_F * maxR2);

        Color densColor = sumColor;// * inv_area; // we want the flux density

        //sumDir *= inv_nc;
        Vector3 avrDir = sumDir.GetNormalize();
        //Vector3 V = hit.incDir * -1.0f;
        
        // TODO:
        //res = PhongDiffuseShade(V, hit.GetTriangleNormal(), sumDir*-1.0f, sumColor, diffuse);
        float nDotL = max(hit.GetTriangleNormal().Dot(avrDir*-1.0f), 0.0f);
        res = diffuse * densColor * nDotL;
        
        nnei = static_cast<float>(neiCount);

        //densColor = sumColor / sumColor.Gray();
        //res = densColor;
        //res.Clamp(0.0f, 1.0f);
        /**/

		return true;
	}
    
                                                                //closest photon index in photon map array
    bool PhotonMap::BinarySearchNearestPht(const Vector3& pos, int & resIndex)
    {
        // ------------------- Binary search the closest photon -------------------
        int zidHit = 0;
		GetZcurveIndex(pos, zidHit);
		int zidL, zidR, zidMid;
		zidL=zidR = 0;
		int leftB = 0; // bound index in photon array
		int rightB = num_pht-1;
		int midB;
        
		Photon phtL = this->LoadPhotonGivenID(leftB);
		Photon phtR = this->LoadPhotonGivenID(rightB);
        
		GetZcurveIndex(phtL.pRay._o, zidL);
		GetZcurveIndex(phtR.pRay._o, zidR);
        
		if(zidHit < zidL || zidHit > zidR) // Don't care if hit point outside photons range
		{
			//res.Set(0.0f);
			return false;
		}
        
		// binary separate searching
		Photon phtM;
		int zidRes;
		//int resIndex; //
        
		//printf("zidHit: %d\n", zidHit);
		while(rightB > leftB+1)
		{
			midB = (leftB+rightB) / 2;
			phtM = this->LoadPhotonGivenID(midB);
            
			GetZcurveIndex(phtM.pRay._o, zidMid);
            
			// compare zid
			//if(zidHit <= zidMid+1 && zidHit >= zidMid-1)
			if(zidHit == zidMid)
			{
				zidRes = zidMid;
				resIndex = midB;
				break;
			}
            
			if(zidHit < zidMid)
			{
				rightB = midB;
				phtR = this->LoadPhotonGivenID(rightB);
				GetZcurveIndex(phtR.pRay._o, zidR);
			}else{
				leftB = midB;
				phtL = this->LoadPhotonGivenID(leftB);
				GetZcurveIndex(phtL.pRay._o, zidL);
			}
            
			zidRes = zidL;
			resIndex = leftB; // by default use left bound photon as result
		}
        
        return true;
		// -------------------------------------------------------------------
    }
    
    bool PhotonMap::BidirZcurveIndexEst(int resIndex, float r2, int maxCount, int maxCheck, HitRecord & hit, Color & sumColor, Vector3 & sumDir, float &omax_r2, int &onei)
    {
        int forward = 0;
		int backward = 1;
		int count = 0;
		int checked = 0;
        
        bool checkF = true;
        bool checkB = true;
		
		Photon phtC; // current loaded photon
		
        float &maxR2 = omax_r2;
        
        int id  = 0;
		while(
              (checkF || checkB) &&
              count < maxCount &&
              checked < maxCheck
              && checked < num_pht-1 //-1 keep it safe?
              )
		{
			if(resIndex + forward < num_pht ||
               resIndex - backward >=0 )
			{
				if(resIndex + forward < num_pht && checkF)
				{
					id = resIndex + forward;
					forward++;
					phtC = this->LoadPhotonGivenID(id);
                    
                    int curZ;
                    GetZcurveIndex(phtC.pRay._o, curZ);
                    if(curZ != resIndex)
                    {
                        checkF = false;
                    }else
                    {
                        float tR2 = (phtC.pRay._o - hit.P).LengthSquare();
                        if(
                           tR2 < r2 &&
                           (phtC.pRay._dir.Dot(hit.GetTriangleNormal()) < 0.0f) )
                        {
                            if(tR2 > maxR2) maxR2 = tR2;
                            // sum++
                            sumColor += phtC.flux;
                            sumDir += phtC.pRay._dir;
                            count++;
                        }

                    }
					   
					checked++;
				}
                
				if(resIndex - backward >= 0 && checkB)
				{
					id = resIndex - backward;
					backward++;
					phtC = this->LoadPhotonGivenID(id);
                    
                    int curZ;
                    GetZcurveIndex(phtC.pRay._o, curZ);
                    if(curZ != resIndex)
                    {
                        checkB = false;
                    }else
                    {
                        float tR2 = (phtC.pRay._o - hit.P).LengthSquare();
                        if(
                           tR2 < r2 &&
                           (phtC.pRay._dir.Dot(hit.GetTriangleNormal()) < 0.0f) )
                        {
                            if(tR2 > maxR2) maxR2 = tR2;
                            // sum++
                            sumColor += phtC.flux;
                            sumDir += phtC.pRay._dir;
                            count++;
                        }
                    }
                    
					checked++;
				}
                
			}else
				break;
            
		}
        
        onei += count;
        return true;
    }
    
    bool PhotonMap::SingleZcurveIndexEst(int resIndex, float r2, int maxCount, int maxCheck, HitRecord & hit, Color & sumColor, Vector3 & sumDir, float &omax_r2, int &onei)
    {
        int forward = 0;
		int backward = 1;
		int count = 0;
		int checked = 0;
        
        bool checkF = true;
        bool checkB = true;
		
		Photon phtC; // current loaded photon
		
        float &maxR2 = omax_r2;
        
        int id  = 0;
		while( count < maxCount &&
               checked < maxCheck
               && checked < num_pht-1 //-1 keep it safe?
              )
		{
			if(resIndex + forward < num_pht ||
               resIndex - backward >=0 )
			{
				if(resIndex + forward < num_pht && checkF)
				{
					id = resIndex + forward;
					forward++;
					phtC = this->LoadPhotonGivenID(id);
                    
                    {
                        float tR2 = (phtC.pRay._o - hit.P).LengthSquare();
                        if(
                           tR2 < r2 &&
                           (phtC.pRay._dir.Dot(hit.GetTriangleNormal()) < 0.0f) )
                        {
                            if(tR2 > maxR2) maxR2 = tR2;
                            // sum++
                            sumColor += phtC.flux;
                            sumDir += phtC.pRay._dir;
                            count++;
                        }
                        
                    }
                    
					checked++;
				}
                
				if(resIndex - backward >= 0 && checkB)
				{
					id = resIndex - backward;
					backward++;
					phtC = this->LoadPhotonGivenID(id);
                    
                    {
                        float tR2 = (phtC.pRay._o - hit.P).LengthSquare();
                        if(
                           tR2 < r2 &&
                           (phtC.pRay._dir.Dot(hit.GetTriangleNormal()) < 0.0f) )
                        {
                            if(tR2 > maxR2) maxR2 = tR2;
                            // sum++
                            sumColor += phtC.flux;
                            sumDir += phtC.pRay._dir;
                            count++;
                        }
                    }
                    
					checked++;
				}
                
			}else
				break;
            
		}
        
        onei += count;
        return true;
    }
    
    //------------------------- Using z-index as grid index ------------------------
    
    bool PhotonMap::LocateFirstPhotonInZindex(const Vector3& pos, int &phtIndex)
    {
        int nearestIndex;
        if(!BinarySearchNearestPht(pos, nearestIndex)){
            //return false; // not in the range of the scene;
        }
        
        Photon pht;
        pht = this->LoadPhotonGivenID(nearestIndex);
        
        int zHit, zNearest;
        GetZcurveIndex(pos, zHit);
        GetZcurveIndex(pht.pRay._o, zNearest);
        
        int curIndex = nearestIndex;
        if(zNearest == zHit)
        {
            // good, within the grid
            int zCur = zNearest;
            
            while(zCur == zHit){
                // find the first photon of that index
                pht = this->LoadPhotonGivenID(curIndex);
                GetZcurveIndex(pht.pRay._o, zCur);
                curIndex--;
                if(curIndex < 0) break;
            }
        }else{
            phtIndex = nearestIndex;
            // TODO:
            return true;
        }
        
        phtIndex = curIndex+1;
        
        return true;
    }
    
    bool PhotonMap::AccumulatePhotonInGrid(int pht1Index, float r2, int maxCount, HitRecord & hit, Color & sumColor, Vector3 & sumDir, float &omax_r2, int &onei)
    {
        int zHead;
        Photon pht = this->LoadPhotonGivenID(pht1Index);
        GetZcurveIndex(pht.pRay._o, zHead);
        
        int curZ = zHead;
        int count = 0;
        int i = pht1Index;
        int checked = 0;
        
        while(curZ == zHead)
        {
            float tR2 = (pht.pRay._o - hit.P).LengthSquare();
            if(
               tR2 < r2 &&
               (pht.pRay._dir.Dot(hit.GetTriangleNormal()) < 0.0f) )
            {
                if(tR2 > omax_r2) omax_r2 = tR2;
                // sum++
                sumColor += pht.flux;
                sumDir += pht.pRay._dir;
                count++;
            }
            i++;
            checked++;
            if(i>num_pht-1) break;
            if(count>maxCount) break;
            //if(checked > 100) break;
            //printf("checked: %d\n", checked);
            pht=this->LoadPhotonGivenID(i);
            GetZcurveIndex(pht.pRay._o, curZ);
        }
        
        onei += count;
        
        return true;
    }
    
    
	//Potential problem: z-curve index doesn't consider negative coordinate
	bool PhotonMap::ZcurveRadEst(HitRecord &hit, const float r, Color &res, float &omax_r2, float &nnei)
	{
		if(s_mat_addr < 0)
			s_mat_addr = GetMaterials();
		int matAddr = s_mat_addr + hit.matID * 25;

		Color diffuse;
		hit.matIllum = loadi(matAddr, 0);
		if(hit.matIllum == 3 || hit.matIllum == 6)
		{
			//printf("Non photon surface\n");
			res.Set(0.0f, 0.0f, 0.0f); // ignore non-photon surface so far
			return true;
		}
		diffuse.LoadFromGM(matAddr+4);
		diffuse.Clamp();

		// TODO: Z-curve sorted and radiance estimation
		if(!sorted){
			printf("thread %d Sorting...\n", _tid);
			this->Sort();
			printf("thread %d Sort done.\n", _tid);
		}

		int resIndex = 0;
        
        if(!BinarySearchNearestPht(hit.P, resIndex))
        {
            res.Set(0.0f);
            return true;
        }
        
        const int oneDim = 3;
        const int size3d = oneDim * oneDim * oneDim;
        int ckIndex[size3d];
        bool toCheck[size3d];
        Vector3 offset;
        
        if(0)
        {
            const int loffset = (oneDim-1)/2;
            float startX = -1.0f;
            float precision = 1.0f;
            for (int i=0; i<DEFAULT_Z_PREC; i++) {
                precision *= 0.1f;
            }
            startX = -precision * static_cast<float>(loffset);
            int curp = 0;
            for(int dx = 0; dx<oneDim; dx++) {
                for (int dy=0; dy<oneDim; dy++) {
                    for (int dz=0; dz<oneDim; dz++) {
                        offset.Set(startX + static_cast<float>(dx)*precision,
                                   startX + static_cast<float>(dy)*precision,
                                   startX + static_cast<float>(dz)*precision);
                        toCheck[curp] = BinarySearchNearestPht(hit.P + offset, ckIndex[curp]);
                        //int t;
                        //toCheck[curp] = LocateFirstPhotonInZindex(hit.P + offset, ckIndex[curp]);
                        curp++;
                    }
                }
            }
        }
        
        //-------- doing multiple cell started point may add a same photon twice in a single estimation pass--------
		// -------------- Estimate irradiance -----------------------
		float r2 = r*r;
        omax_r2 = 0.0f;
        float &maxR2 = omax_r2;
		
        Color sumColor = 0.0f;
		Vector3 sumDir = 0.0f;
        
		int count = 0;
        int maxCount = PHT_COUNT / size3d;
        int maxCheck = PHT_CHECKED / size3d;
        
        if(0)
        for (int i=0; i<size3d; i++) {
            if(toCheck[i]){
                //BidirZcurveIndexEst(ckIndex[i], r2, maxCount, maxCheck, hit, sumColor, sumDir, maxR2, count);
                AccumulatePhotonInGrid(ckIndex[i], r2, PHT_COUNT, hit, sumColor, sumDir, maxR2, count);
            }
        }
        SingleZcurveIndexEst(resIndex, r2, PHT_COUNT, PHT_CHECKED, hit, sumColor, sumDir, omax_r2, count);
        
		if(count < 8){
			//res.Set(0.0f);
			//return true;
		}
		if(maxR2 < 0.01f || maxR2 > r2)
            maxR2 = r2;
        // maxR2 = 10.0f;
		float area = PI_F * maxR2;
		Color densColor = sumColor;// / area;
		Vector3 avrDir = sumDir.GetNormalize();

		float nDotL = max(hit.GetTriangleNormal().Dot(avrDir*-1.0f), 0.0f);
		res = diffuse * densColor * nDotL;
        res.ClampNeg();
        nnei = static_cast<float>(count);
        
        // ----------------------------------------------------
		//densColor = sumColor / sumColor.Gray();
		//res = densColor;
		//res = sumColor;
		//if(res.R > 1.0f || res.B > 1.0f || res.G > 1.0f)
		//{
			//res.Set(0.2f, 0.2f, 1.0f);
		//}

		return true;
	}

	bool PhotonMap::RadianceEstimate(HitRecord &hit, const float r, Color & res, float &omax_r2, float &nnei)
	{
		if(pm_s_addr < 0)
			pm_s_addr = GetPhotonMapAddr();

		if(pm_s_addr < 0)
			return false; // photon map adress incorrect

		// Brute force neighbors search
		if(num_pht < 0)
			num_pht = loadi(pm_s_addr, 0);

		if(s_mat_addr < 0)
		    s_mat_addr = GetMaterials();

        /**
		if(!sorted)
		{
			printf("thread %d whatever Sorting...\n", _tid);
			this->Sort();
			printf("thread %d whatever Sort done.\n", _tid);
		}
        // */
        
		int matAddr = s_mat_addr + hit.matID * 25;
        hit.transIndex = loadf(matAddr, 16);

		//return this->ZcurveRadEst(hit, r, res, omax_r2, nnei);

		return this->BruteForceRadEst(hit, r, res, omax_r2, nnei);
	}
    
    void PhotonMap::Reset()
    {
        if(_tid == 0)
        {
            int cur_pht_offset = atomicadd(REG_PHT_ST_ID, -num_pht);
            //printf("reset number of photons: %d\n", cur_pht_offset);
            storei(0, pm_s_addr, 0);
        }
        sorted = false;
        barrier(REG_BARRIER);
        num_pht = 0;
        //printf("thread %d know %d photons.\n", _tid, num_pht);
    }
}
